import sys
import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.utils.data import Subset
from sklearn.model_selection import KFold

# user-defined package
import config
from utils import ddp_utils
from models import VGGNet
from datasets import FoodDataset
from networks.datasets.FoodDataset import MixupDataset


def calc_accuracy(output, target):
    return output.argmax(dim=-1).eq(target).float().mean()

def get_model():
    return VGGNet.get_model()

def get_criterion():
    return nn.CrossEntropyLoss()

def get_optimizer(parameters):
    optimizer = optim.Adam(parameters, lr=config.VGGNetConfig.learning_rate, weight_decay=config.VGGNetConfig.weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=config.VGGNetConfig.scheduler_factor, patience=config.VGGNetConfig.scheduler_patience)
    return optimizer, scheduler

def get_dataset(train_ids, valid_ids):
    train_sub = Subset(config.VGGNetConfig.full_dataset, train_ids)
    valid_sub = Subset(config.VGGNetConfig.full_dataset, valid_ids)

    train_tfm, test_tfm = FoodDataset.get_transforms(train_sub)
    train_sub.dataset.transforms = train_tfm
    valid_sub.dataset.transforms = test_tfm

    train_dataset = MixupDataset(train_sub)
    valid_dataset = valid_sub
    return train_dataset, valid_dataset

if __name__ == "__main__":

    world_size = torch.cuda.device_count()
    writer_path = config.VGGNetConfig.writer_path
    log_path = config.VGGNetConfig.log_path
    model_path = config.VGGNetConfig.checkpoint_path
    epochs = config.VGGNetConfig.epochs
    batch_size = config.VGGNetConfig.batch_size
    patience = config.VGGNetConfig.valid_patience
    kFold = KFold(n_splits=config.VGGNetConfig.kFold, shuffle=True)
    dataset_size = config.VGGNetConfig.dataset_size

    try:
        mp.spawn(ddp_utils.train_worker, nprocs=world_size, args=(
            world_size, get_model, get_criterion, get_optimizer, get_dataset,
            calc_accuracy, epochs, batch_size, patience, False, kFold, dataset_size,
            writer_path, log_path, model_path
        ))
    except Exception as e:
        print(f"Main process exception: {e}")
    except KeyboardInterrupt:
        print(f"KeyboardInterrupt detected. Terminating process group")
    finally:
        for p in mp.active_children():
            print(f"Terminating process {p.pid}")
            p.terminate()
        sys.exit(1)


